# -*- coding: utf-8 -*-
"""
@author:xuyuntao
@time:2021/3/1:9:12
@email:xuyuntao@189.cn
"""
import numpy
import numpy.fft as nfft

"""计算各种瞬时特征，注意结尾带_m的函数必须输入二维矩阵，否则计算结果错误
所有输入signal参数均假设一行为一个信号的采样序列，可以有多行，也可以只有一行
末尾带_avg的函数针对同时输入多个信号的情况，求其平均值"""

def R(signal:numpy.ndarray):
    oneDim=False
    if len(signal.shape)==1:
        oneDim=True
        signal=signal.reshape([1,-1])
    if numpy.iscomplex(signal).any():
        envelope=numpy.abs(signal)
        sigma2=numpy.var(envelope**2,axis=1)
        mu=numpy.mean(envelope**2,axis=1)
        r=sigma2/(mu**2)
        if oneDim:
            r = r[0]
        return r
    else:
        raise ValueError("输入信号必须为复调制信号")

def A_cn_m(signal:numpy.ndarray):
    if numpy.iscomplex(signal).any():
        amplitude=numpy.abs(signal.real)
    else:
        amplitude = numpy.abs(signal)
    m_a=numpy.mean(amplitude,axis=1).reshape([-1,1]).repeat(amplitude.shape[1],axis=1)
    a_n=amplitude/m_a
    a_cn=a_n-1
    return a_cn

def Gamma_max(signal:numpy.ndarray):
    oneDim=False
    if len(signal.shape)==1:
        oneDim=True
        signal=signal.reshape([1,-1])
    gamma_max=numpy.abs(nfft.fft(A_cn_m(signal),axis=1)).max(axis=1)**2/signal.shape[1]
    if oneDim:
        gamma_max = gamma_max[0]
    return gamma_max

def Sigma_aa(signal:numpy.ndarray):
    oneDim=False
    if len(signal.shape)==1:
        oneDim=True
        signal=signal.reshape([1,-1])
    a_cn=A_cn_m(signal)
    sigma_aa=numpy.sqrt(numpy.mean(a_cn**2,axis=1)-numpy.mean(numpy.abs(a_cn),axis=1)**2)
    if oneDim:
        sigma_aa = sigma_aa[0]
    return sigma_aa

def Phi_NL_m(signal:numpy.ndarray,timeSeries:numpy.ndarray,carrierFreq:int,
             noiseRad:float=0.1*numpy.pi):
    """考虑两种情况：一是像源程序中一样hilbert变换出虚部，二是直接利用复信号
    绘图看结果后选择"""
    if signal.shape==timeSeries.shape:
        if numpy.iscomplex(signal).any():
            phase=numpy.arctan2(signal.imag,signal.real)
            # fix#me:验证是否正确，pi/2附近如何划分
            phase_tmp=numpy.zeros([phase.shape[0],phase.shape[1]+1],dtype=numpy.float64)
            phase_tmp[:,:-1]=phase
            phase_diff=numpy.diff(phase_tmp,axis=1)
            phase_correct_sep=numpy.zeros(phase.shape,dtype=numpy.float64)
            phase_correct_sep[phase_diff>numpy.pi]=-2*numpy.pi
            phase_correct_sep[phase_diff<-1*numpy.pi]=2*numpy.pi
            phase_correct=numpy.cumsum(phase_correct_sep,axis=1)
            phi_NL=phase + phase_correct + (2 * numpy.pi * timeSeries * carrierFreq)
            return numpy.mod(phi_NL+noiseRad, 2*numpy.pi)-noiseRad
            # fix#me:验证线性量
        else:
            raise ValueError("输入信号必须为复调制信号")
    else:
        raise ValueError("信号序列和时间序列大小要保持一致")

def Sigma_ap(signal:numpy.ndarray, timeSeries:numpy.ndarray, carrierFreq:int,
             a_t:float=1.0, noiseRad:float=0.1*numpy.pi):
    oneDim=False
    if len(signal.shape)==1:
        oneDim=True
        signal=signal.reshape([1,-1])
    if signal.shape==timeSeries.shape:
        if numpy.iscomplex(signal).any():
            amplitude = numpy.abs(signal.real)
            m_a=numpy.mean(amplitude,axis=1).reshape([-1,1]).repeat(amplitude.shape[1],axis=1)
            a_n = amplitude / m_a
            judgeArray=a_n>a_t
            C=judgeArray.sum(axis=1,dtype=numpy.int32)
            C[C==0]=1
            phi_NL=Phi_NL_m(signal,timeSeries,carrierFreq,noiseRad)
            phi_NL[~judgeArray]=0
            sigma_ap_square=(numpy.sum(phi_NL**2,axis=1))/C-\
                            (numpy.sum(numpy.abs(phi_NL),axis=1)/C)**2
            sigma_ap_square[C==0]=1
            sigma_ap_square[sigma_ap_square<0]=0
            sigma_ap=numpy.sqrt(sigma_ap_square)
            if oneDim:
                sigma_ap = sigma_ap[0]
            return sigma_ap
        else:
            raise ValueError("输入信号必须为复调制信号")
    else:
        raise ValueError("信号序列和时间序列大小要保持一致")

def F_NL_m(signal:numpy.ndarray,timeSeries:numpy.ndarray,carrierFreq:int,
           noiseRad:float=0.1*numpy.pi):
    if signal.shape == timeSeries.shape:
        if numpy.iscomplex(signal).any():
            phi_NL = Phi_NL_m(signal, timeSeries, carrierFreq,noiseRad)
            phi_NL_tmp=numpy.zeros([phi_NL.shape[0],phi_NL.shape[1]+1])
            phi_NL_tmp[:,1:]=phi_NL
            f_NL=numpy.diff(phi_NL_tmp,axis=1)
            return f_NL
        else:
            raise ValueError("输入信号必须为复调制信号")
    else:
        raise ValueError("信号序列和时间序列大小要保持一致")

def Sigma_af(signal:numpy.ndarray, timeSeries:numpy.ndarray, carrierFreq:int,
             a_t:float=1.0, noiseRad:float=0.1*numpy.pi):
    oneDim=False
    if len(signal.shape)==1:
        oneDim=True
        signal=signal.reshape([1,-1])
    if signal.shape == timeSeries.shape:
        if numpy.iscomplex(signal).any():
            amplitude = numpy.abs(signal.real)
            m_a = numpy.mean(amplitude,axis=1).reshape([-1,1]).repeat(amplitude.shape[1],axis=1)
            a_n = amplitude / m_a
            judgeArray=a_n>a_t
            C = judgeArray.sum(axis=1, dtype=numpy.int32)
            C[C == 0] = 1
            f_NL=F_NL_m(signal, timeSeries, carrierFreq,noiseRad)
            f_NL[~judgeArray] = 0
            sigma_af_square=(numpy.sum(f_NL**2,axis=1))/C-\
                            (numpy.sum(numpy.abs(f_NL),axis=1)/C)**2
            sigma_af_square[C==0]=1
            sigma_af_square[sigma_af_square<0]=0
            sigma_af = numpy.sqrt(sigma_af_square)
            if oneDim:
                sigma_af = sigma_af[0]
            return sigma_af
        else:
            raise ValueError("输入信号必须为复调制信号")
    else:
        raise ValueError("信号序列和时间序列大小要保持一致")

def R_avg(signal:numpy.ndarray):
    return R(signal).mean()

def Gamma_max_avg(signal:numpy.ndarray):
    return Gamma_max(signal).mean()

def Sigma_aa_avg(signal:numpy.ndarray):
    return Sigma_aa(signal).mean()

def Sigma_ap_avg(signal:numpy.ndarray, timeSeries:numpy.ndarray, carrierFreq:int,
                 a_t:float=1.0, noiseRad:float=0.1*numpy.pi):
    return Sigma_ap(signal, timeSeries, carrierFreq, a_t, noiseRad).mean()

def Sigma_af_avg(signal:numpy.ndarray, timeSeries:numpy.ndarray, carrierFreq:int,
                 a_t:float=1.0, noiseRad:float=0.1*numpy.pi):
    return Sigma_af(signal, timeSeries, carrierFreq, a_t, noiseRad).mean()